package io.github.wanggit.antrpc.client.zk.lb;

import io.github.wanggit.antrpc.client.zk.zknode.NodeHostEntity;
import io.github.wanggit.antrpc.commons.bean.Host;
import io.github.wanggit.antrpc.commons.bean.RegisterBeanMethod;
import io.github.wanggit.antrpc.commons.utils.JsonUtils;
import lombok.extern.slf4j.Slf4j;

import java.util.*;
import java.util.concurrent.CopyOnWriteArrayList;

@Slf4j
abstract class AbstractLoadBalancer implements ILoadBalancer {

    CopyOnWriteArrayList<NodeHostEntity> entities = new CopyOnWriteArrayList<>();

    private Map<String, Integer> weights;

    /**
     * 统计IP的使用频率
     *
     * @param entity host
     */
    void stats(Host entity) {
        String hostInfo = entity.getHostInfo();
        HostLogHolder.getInstance().log(hostInfo);
    }

    void copyEntitiesTo(AbstractLoadBalancer loadBalancer) {
        loadBalancer.setEntities(entities);
    }

    private void setEntities(CopyOnWriteArrayList<NodeHostEntity> entities) {
        this.entities = entities;
    }

    @Override
    public boolean isEmpty() {
        return entities.isEmpty();
    }

    @Override
    public void setWeights(Map<String, Integer> weights) {
        if (null == weights || weights.isEmpty()) {
            return;
        }
        this.weights = weights;
        if (entities.isEmpty()) {
            if (log.isWarnEnabled()) {
                log.warn(
                        "Set weights to loadBalancer, but entities is empty. will fix it when addHost() called.");
            }
            return;
        }
        List<NodeHostEntity> hostEntities = new ArrayList<>(entities.size());
        Set<NodeHostEntity> hostEntitySet = new HashSet<>(entities);
        if (log.isDebugEnabled()) {
            log.debug(
                    "weights = "
                            + JsonUtils.toJsonString(weights)
                            + ", current hosts = "
                            + JsonUtils.toJsonString(hostEntitySet));
        }
        for (NodeHostEntity hostEntity : hostEntitySet) {
            Integer weight = weights.getOrDefault(hostEntity.getHostInfo(), 1);
            for (int i = 0; i < weight; i++) {
                hostEntities.add(hostEntity);
            }
        }
        entities = new CopyOnWriteArrayList<>(hostEntities);
    }

    @Override
    public void addHost(NodeHostEntity host) {
        if (!entities.contains(host)) {
            if (null == weights || weights.isEmpty()) {
                entities.add(host);
            } else {
                Integer weight = weights.getOrDefault(host.getHostInfo(), 1);
                if (log.isInfoEnabled()) {
                    log.info("add " + host.getHostInfo() + " to loadBalancer, weight is " + weight);
                }
                for (int i = 0; i < weight; i++) {
                    entities.add(host);
                }
            }
        } else {
            if (log.isInfoEnabled()) {
                log.info("add " + host.getHostInfo() + " to loadBalancer, but is existed.");
            }
        }
    }

    @Override
    public void updateHost(NodeHostEntity host) {
        if (!entities.contains(host)) {
            addHost(host);
        } else {
            // 检查权重是否正确
            if (null != weights) {
                Integer weight = weights.getOrDefault(host.getHostInfo(), 1);
                long count =
                        entities.stream()
                                .filter(it -> Objects.equals(host.getHostInfo(), it.getHostInfo()))
                                .count();
                if (weight != count) {
                    if (log.isInfoEnabled()) {
                        log.info(
                                "check the "
                                        + host
                                        + ", expected weight = "
                                        + weight
                                        + ", but actual = "
                                        + count
                                        + ", It will adjust automatically.");
                    }
                    deleteHost(host);
                    addHost(host);
                }
            }
        }
    }

    @Override
    public void deleteHost(NodeHostEntity host) {
        boolean result;
        do {
            result = entities.remove(host);
        } while (result);
    }

    @Override
    public List<NodeHostEntity> getAllNodeHostEntities() {
        return new ArrayList<>(entities);
    }

    @Override
    public void updateHostRateLimit(Host host, RegisterBeanMethod registerBeanMethod) {
        entities.forEach(
                it -> {
                    if (Objects.equals(it.getHostInfo(), host.getHostInfo())) {
                        Map<String, RegisterBeanMethod> methodMap = it.getMethodMap();
                        if (null != methodMap) {
                            RegisterBeanMethod oldRegisterBeanMethod =
                                    methodMap.get(registerBeanMethod.toString());
                            if (registerBeanMethod.getLimit() > 0
                                    && registerBeanMethod.getDurationInSeconds() > 0) {
                                oldRegisterBeanMethod.setLimit(registerBeanMethod.getLimit());
                                oldRegisterBeanMethod.setDurationInSeconds(
                                        registerBeanMethod.getDurationInSeconds());
                            }
                        }
                    }
                });
    }
}
